package org.west.sky.scripture.ratelimiter;

import java.time.LocalTime;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author chz
 * @date 2022/3/18
 * @description: 滑动窗口限流算法
 */
public class SlideWindowRateLimiter {

    /**
     * 阈值
     */
    private int qps = 2;
    /**
     * 时间窗口总大小（毫秒）
     */
    private long windowSize = 1000;
    /**
     * 多少个子窗口
     */
    private Integer windowCount = 10;
    /**
     * 窗口列表
     */
    private WindowInfo[] windowArray = new WindowInfo[windowCount];

    public SlideWindowRateLimiter(int qps) {
        this.qps = qps;
        long currentTimeMillis = System.currentTimeMillis();
        for (int i = 0; i < windowArray.length; i++) {
            windowArray[i] = new WindowInfo(currentTimeMillis, new AtomicInteger(0));
        }
    }

    /**
     * 1. 计算当前时间窗口
     * 2. 更新当前窗口计数 & 重置过期窗口计数
     * 3. 当前 QPS 是否超过限制
     *
     * @return
     */
    public synchronized boolean tryAcquire() {
        long currentTimeMillis = System.currentTimeMillis();
        //计算当前时间窗口，先算出当中1s中的所在毫秒数除以时间刻度得出所在窗口
        int currentIndex = (int) ((currentTimeMillis % windowSize) / (windowSize / windowCount));
        int sum = 0;
        //更新当前窗口计数、重置过期窗口计数
        for (int i = 0; i < windowArray.length; i++) {
            WindowInfo windowInfo = windowArray[i];
            //重置过期窗口计数
            if (currentTimeMillis - windowInfo.getTime() > windowSize) {
                windowInfo.getNumber().set(0);
                windowInfo.setTime(currentTimeMillis);
            }
            //更新当前窗口计数
            if (currentIndex == i && windowInfo.getNumber().get() < qps) {
                windowInfo.getNumber().incrementAndGet();
            }
            sum = sum + windowInfo.getNumber().get();
        }
        return sum <= qps;
    }

    /**
     * 设置 QPS 为 2,测试次数 20 次，每次间隔 300 毫秒，预计成功次数在 12 次左右
     *
     * @param args
     */
    public static void main(String[] args) throws InterruptedException {
        int qps = 2, count = 20, sleep = 300, success = count * sleep / 1000 * qps;
        System.out.println(String.format("当前QPS限制为:%d,当前测试次数为:%d,间隔:%dms,预计成功:%d次", qps, count, sleep, success));
        success = 0;
        SlideWindowRateLimiter slideWindowRateLimiter = new SlideWindowRateLimiter(2);
        for (int i = 0; i < count; i++) {
            Thread.sleep(sleep);
            if (slideWindowRateLimiter.tryAcquire()) {
                success++;
                if (success % qps == 0) {
                    System.out.println(LocalTime.now() + ": success,");
                } else {
                    System.out.print(LocalTime.now() + ": success,");
                }
            } else {
                System.out.println(LocalTime.now() + ": fail");
            }
        }
        System.out.println();
        System.out.println("实际测试成功次数:" + success);
    }

    /**
     * 窗口信息
     */
    private class WindowInfo {
        /**
         * 窗口开始时间
         */
        private Long time;
        /**
         * 计数器
         */
        private AtomicInteger number;

        public WindowInfo(long time, AtomicInteger number) {
            this.time = time;
            this.number = number;
        }

        public Long getTime() {
            return time;
        }

        public void setTime(Long time) {
            this.time = time;
        }

        public AtomicInteger getNumber() {
            return number;
        }

        public void setNumber(AtomicInteger number) {
            this.number = number;
        }
    }

}
